package com.kujie.ai.config;

import com.kujie.ai.advisor.ReasoningContentAdvisor;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.MessageChatMemoryAdvisor;
import org.springframework.ai.chat.client.advisor.PromptChatMemoryAdvisor;
import org.springframework.ai.chat.client.advisor.SimpleLoggerAdvisor;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.memory.ChatMemoryRepository;
import org.springframework.ai.chat.memory.MessageWindowChatMemory;
import org.springframework.ai.chat.memory.repository.jdbc.JdbcChatMemoryRepository;
import org.springframework.ai.chat.memory.repository.jdbc.MysqlChatMemoryRepositoryDialect;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.jdbc.core.JdbcTemplate;
import org.springframework.jdbc.support.JdbcTransactionManager;

import javax.sql.DataSource;
import java.util.List;

/**
 * 客户端 配置
 *
 * @author check
 * @date 2025-07-24
 */
@Configuration
public class ChatClientConfig {
    @Autowired
    private DataSource dataSource;

//    @Autowired
//    private RetrievalAugmentationAdvisor retrievalAugmentationAdvisor;

    @Autowired
    private SimpleLoggerAdvisor simpleLoggerAdvisor;
//
//    @Autowired
//    private MilvusVectorStore milvusVectorStore;

    @Bean
    public ChatMemory mysqlChatMemoryRepository() {
        JdbcTemplate jdbcTemplate = new JdbcTemplate(dataSource);
        ChatMemoryRepository chatMemoryRepository = JdbcChatMemoryRepository.builder()
                .jdbcTemplate(jdbcTemplate).dialect(new MysqlChatMemoryRepositoryDialect())
                .transactionManager(new JdbcTransactionManager(dataSource)).build();

        // 清除上下文记忆
        List<String> conversationIds = chatMemoryRepository.findConversationIds();
        for (String conversationId : conversationIds) {
            chatMemoryRepository.deleteByConversationId(conversationId);
        }

        return MessageWindowChatMemory.builder().chatMemoryRepository(chatMemoryRepository).maxMessages(5).build();
    }

    @Bean
    public ChatClient buildChatClient(ChatClient.Builder ollamaChatModel) {
        return ollamaChatModel.build();
    }

    @Bean
    public ChatClient buildCurrentChatClient(ChatClient.Builder ollamaChatModel, ChatMemory chatMemory) {
        return ollamaChatModel
                .defaultAdvisors(
                        MessageChatMemoryAdvisor.builder(chatMemory).build(),
//                        new QuestionAnswerAdvisor(milvusVectorStore),
                        PromptChatMemoryAdvisor.builder(chatMemory).build(),
//                        retrievalAugmentationAdvisor,
                        simpleLoggerAdvisor,
                        new ReasoningContentAdvisor(0))
                .build();
    }
}
